import cv2
import os
import random

def laplacian_pyramid_blending(img1, img2, mask):
    # 生成高斯金字塔
    G1, G2, Gm = img1.copy(), img2.copy(), mask.copy()
    gp1, gp2, gp_m = [G1], [G2], [Gm]

    for i in range(6):
        G1 = cv2.pyrDown(G1)
        G2 = cv2.pyrDown(G2)
        Gm = cv2.pyrDown(Gm)
        gp1.append(G1)
        gp2.append(G2)
        gp_m.append(Gm)

    # 生成拉普拉斯金字塔
    lp1, lp2 = [gp1[-1]], [gp2[-1]]
    for i in range(5, 0, -1):
        L1 = cv2.subtract(gp1[i-1], cv2.pyrUp(gp1[i]))
        L2 = cv2.subtract(gp2[i-1], cv2.pyrUp(gp2[i]))
        lp1.append(L1)
        lp2.append(L2)

    # 融合金字塔
    LS = []
    for l1, l2, gm in zip(lp1, lp2, gp_m):
        LS.append(l1 * gm + l2 * (1.0 - gm))

    # 重建图像
    result = LS[0]
    for i in range(1, 6):
        result = cv2.pyrUp(result)
        result = cv2.add(result, LS[i])

    return result


def process_images(folder_a, folder_b, mask_folder, output_folder):
    """
    遍历前景图像，并使用对应的掩膜进行泊松融合。
    :param folder_a: 前景图像文件夹路径
    :param folder_b: 背景图像文件夹路径
    :param mask_folder: 掩膜文件夹路径
    :param output_folder: 输出图像文件夹路径
    """
    # 确保输出文件夹存在
    os.makedirs(output_folder, exist_ok=True)

    # 获取文件夹中的所有文件名
    files_a = sorted(os.listdir(folder_a))
    files_b = sorted(os.listdir(folder_b))

    if not files_a or not files_b:
        print("前景或背景文件夹为空，无法进行融合。")
        return

    # 遍历前景文件，并使用对应的掩膜文件
    for i, file_a in enumerate(files_a):
        file_b = random.choice(files_b)  # 随机选择一个背景图像
        mask_path = os.path.join(mask_folder, file_a)  # 掩膜文件路径

        # 加载前景、背景和掩膜图像
        foreground = cv2.imread(os.path.join(folder_a, file_a))
        background = cv2.imread(os.path.join(folder_b, file_b))
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        if foreground is None or background is None:
            print(f"无法加载图像：{file_a} 或 {file_b}")
            continue

        if mask is None:
            print(f"掩膜文件缺失，跳过：{file_a}")
            continue

        # 确保前景、背景和掩膜大小一致
        background = cv2.resize(background, (foreground.shape[1], foreground.shape[0]))
        mask = cv2.resize(mask, (foreground.shape[1], foreground.shape[0]))

        # 确保掩膜是二值图像
        _, mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)

        # 定义融合区域的中心点（背景图像的中心）
        center = (background.shape[1] // 2, background.shape[0] // 2)

        # 使用 高斯金字塔融合
        result = laplacian_pyramid_blending(foreground, background, mask)

        # 保存结果图像
        output_path = os.path.join(output_folder, f"Img_{i}.png")
        cv2.imwrite(output_path, result)

        print(f"已保存融合图像：{output_path}")

# 指定文件夹路径
folder_a = "rawA"  # 前景图像文件夹路径
folder_b = "rawB"  # 背景图像文件夹路径
mask_folder = "msk"  # 掩膜文件夹路径
output_folder = "C2"  # 输出图像文件夹路径

# 执行融合
process_images(folder_a, folder_b, mask_folder, output_folder)
